package com.algorithm.learning.base.图;

import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.Stack;

/**
 * @program: algorithm-learning
 * @description:
 * @author: YuKai Fan
 * @create: 2024/8/29 22:19
 **/
public class 最小生成树算法之Kruskal {

    /**
     * 所谓最小生成树，即：
     * 找到图中所有点都可以连接在一起(不破坏连通性)，而且边的权值之和最小
     *
     *  Kruskal -> K算法：需要理解并查集
     * @see 并查集
     *
     *  1、先把所有的边，根据权值从小到大排序
     *  2、遍历所有边，如果边的两个节点不在同一个集合中，则将from和to点设置为一个集合，然后把边加入到结果中
     *  3、如果所有的边都遍历完了，那么结果就是最小生成树
     *
     *  根据上面的步骤可以看到，Kruskal算法是利用了并查集的结构
     * @param graph
     */
    public static Set<Edge> kruskalMST(Graph graph) {
        UnionFind unionFind = new UnionFind();
        unionFind.makeSets(graph.nodes.values());
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());

        for (Edge edge : graph.edges) {
            priorityQueue.add(edge);
        }
        Set<Edge> result = new HashSet<>();
        while (!priorityQueue.isEmpty()) {
            Edge edge = priorityQueue.poll();
            if (!unionFind.isSameSet(edge.from, edge.to)) {
                result.add(edge);
                unionFind.union(edge.from, edge.to);
            }
        }
        return result;
    }

    public static class EdgeComparator implements Comparator<Edge> {

        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;
        }

    }

    public static class UnionFind {
        private Map<Node, Node> parentMap;
        private Map<Node, Integer> sizeMap;

        public UnionFind() {
            this.parentMap = new HashMap<>();
            this.sizeMap = new HashMap<>();
        }

        public void makeSets(Collection<Node> nodes) {
            parentMap.clear();
            sizeMap.clear();
            for (Node node : nodes) {
                parentMap.put(node, node);
                sizeMap.put(node, 1);
            }
        }

        private Node findParent(Node n) {
            Stack<Node> path = new Stack<>();
            while (n != parentMap.get(n)) {
                path.push(n);
                n = parentMap.get(n);
            }
            while (!path.isEmpty()) {
                parentMap.put(path.pop(), n);
            }
            return n;
        }

        /**
         * 判断是否同一集合
         * @param a
         * @param b
         * @return
         */
        public boolean isSameSet(Node a, Node b) {
            if (a == null || b == null) {
                return false;
            }
            return findParent(a) == findParent(b);
        }

        public void union(Node a, Node b) {
            if (a == null || b == null) {
                return;
            }
            Node aParent = findParent(a);
            Node bParent = findParent(b);
            if (aParent != bParent) {
                int aSetSize = sizeMap.get(aParent);
                int bSetSize = sizeMap.get(bParent);
                if (aSetSize >= bSetSize) {
                    parentMap.put(bParent, aParent);
                    sizeMap.put(aParent, aSetSize + bSetSize);
                    sizeMap.remove(bParent);
                } else {
                    parentMap.put(aParent, bParent);
                    sizeMap.put(bParent, aSetSize + bSetSize);
                    sizeMap.remove(aParent);
                }
            }
        }
    }
}
